#include <stdio.h>  
#include <stdint.h>  
#include <stdlib.h>  
#include <string.h>  
#include "gmssl/sm4.h" // 确保这个路径正确指向 GmSSL 库中的 sm4.h  
  
// 函数声明  
void handleErrors(const char *message);  
void sm4_encrypt_file(const char *key_file, const char *input_file, const char *output_file);  
void sm4_decrypt_file(const char *key_file, const char *input_file, const char *output_file);  
  
int main(int argc, char *argv[]) {  
    if (argc != 5) {  
        fprintf(stderr, "Usage: %s <encrypt|decrypt> <key_file> <input_file> <output_file>\n", argv[0]);  
        return EXIT_FAILURE;  
    }  
  
    if (strcmp(argv[1], "encrypt") == 0) {  
        sm4_encrypt_file(argv[2], argv[3], argv[4]);  
    } else if (strcmp(argv[1], "decrypt") == 0) {  
        sm4_decrypt_file(argv[2], argv[3], argv[4]);  
    } else {  
        fprintf(stderr, "Invalid operation: %s\n", argv[1]);  
        return EXIT_FAILURE;  
    }  
  
    return EXIT_SUCCESS;  
}  
  
void sm4_encrypt_file(const char *key_file, const char *input_file, const char *output_file) {  
    FILE *kf, *inf, *outf;  
    SM4_KEY sm4_key;  
    uint8_t key[SM4_KEY_SIZE];  
    uint8_t input[SM4_BLOCK_SIZE];  
    uint8_t output[SM4_BLOCK_SIZE];  
    size_t bytes_read;  
  
    // 打开密钥文件并读取密钥  
    kf = fopen(key_file, "rb");  
    if (!kf) {  
        handleErrors("Failed to open key file");  
    }  
    if (fread(key, 1, SM4_KEY_SIZE, kf) != SM4_KEY_SIZE) {  
        fclose(kf);  
        handleErrors("Failed to read key from file");  
    }  
    fclose(kf);  
  
    // 设置 SM4 加密密钥  
    sm4_set_encrypt_key(&sm4_key, key);  
  
    // 打开输入和输出文件  
    inf = fopen(input_file, "rb");  
    if (!inf) {  
        handleErrors("Failed to open input file");  
    }  
    outf = fopen(output_file, "wb");  
    if (!outf) {  
        fclose(inf);  
        handleErrors("Failed to open output file");  
    }  
  
    // 逐块读取输入文件并进行加密  
    while ((bytes_read = fread(input, 1, SM4_BLOCK_SIZE, inf)) > 0) {  
        // 如果读取的块不足一个完整的块，进行填充（对于加密通常是必要的）  
        if (bytes_read < SM4_BLOCK_SIZE) {  
            memset(input + bytes_read, 0, SM4_BLOCK_SIZE - bytes_read);  
        }  
        sm4_encrypt(&sm4_key, input, output);  
        fwrite(output, 1, SM4_BLOCK_SIZE, outf);  
    }  
  
    fclose(inf);  
    fclose(outf);  
    printf("Encryption completed successfully.\n");  
}  
  
void sm4_decrypt_file(const char *key_file, const char *input_file, const char *output_file) {  
    FILE *kf, *inf, *outf;  
    SM4_KEY sm4_key;  
    uint8_t key[SM4_KEY_SIZE];  
    uint8_t input[SM4_BLOCK_SIZE];  
    uint8_t output[SM4_BLOCK_SIZE];  
    size_t bytes_read;  
  
    // 打开密钥文件并读取密钥  
    kf = fopen(key_file, "rb");  
    if (!kf) {  
        handleErrors("Failed to open key file");  
    }  
    if (fread(key, 1, SM4_KEY_SIZE, kf) != SM4_KEY_SIZE) {  
        fclose(kf);  
        handleErrors("Failed to read key from file");  
    }  
    fclose(kf);  
  
    // 设置 SM4 解密密钥（与加密使用相同的密钥）  
    sm4_set_decrypt_key(&sm4_key, key);  
  
    // 打开输入和输出文件  
    inf = fopen(input_file, "rb");  
    if (!inf) {  
        handleErrors("Failed to open input file");  
    }  
    outf = fopen(output_file, "wb");  
    if (!outf) {  
        fclose(inf);  
        handleErrors("Failed to open output file");  
    }  
  
    // 逐块读取输入文件并进行解密  
    while ((bytes_read = fread(input, 1, SM4_BLOCK_SIZE, inf)) > 0) {  
        sm4_encrypt(&sm4_key, input, output);  
        fwrite(output, 1, bytes_read, outf); // 注意：这里写入的是实际读取的字节数，而不是固定的块大小  
    }  
  
    fclose(inf);  
    fclose(outf);  
    printf("Decryption completed successfully.\n");  
}  
  
void handleErrors(const char *message) {  
    fprintf(stderr, "Error: %s\n", message);  
    exit(EXIT_FAILURE);  
}
